import pandas as pd 
from mlxtend.preprocessing  import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

df = pd.read_excel('./associationRules/餐厅数据.xlsx')
print(df.head)

trsations = df['菜品'].str.split(',').to_list()

te = TransactionEncoder()
te_ary = te.fit(trsations).transform(trsations)
df_enecoded = pd.DataFrame(te_ary, columns=te.columns_)

#使用apriori进行分析
frquent_itemsets = apriori(df_enecoded, min_support=0.1, use_colnames=True)

frquent_itemsets.sort_values(by='support', ascending=False, inplace=True)

print(frquent_itemsets [frquent_itemsets .itemsets.apply(lambda x :len(x) == 2)])

#生成关联规则
rules = association_rules(frquent_itemsets, metric="confidence", min_threshold=0.1)

# 筛选强关联规则
effctive = rules[
    (rules['lift'] > 1)& (rules['confidence'] > 0.1)
].sort_values(by=['lift', 'confidence'], ascending=False) 

print(effctive[['antecedents', 'consequents','support','confidence', 'lift']])

import matplotlib.pyplot as plt
import networkx as nx

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
G = nx.DiGraph()
for _, row in effctive.iterrows():
    G.add_edge(','.join(list(row['antecedents'])),
                ','.join(list(row['antecedents'])),
                weight=row['lift'])

plt.figure(figsize = (12, 8))
pos = nx.spring_layout(G)
nx.draw(G, pos,
        with_lables=True,
        edge_color=[G[u][v]['weight'] for u , v in G.edges()],
        width=2.0,
        edge_cmap=plt.cm.Blues
        )               
plt.show()

